/*
  Author: Chen Haodong (shinkunyachen@outlook.com)
  -----
  Copyright (c)2020 - present Chen Haodong. All Rights Reserved.
 */

#ifndef BRANCHCHECKER_LCFG_H
#define BRANCHCHECKER_LCFG_H

#include <clang/AST/Decl.h>
#include <clang/AST/Stmt.h>
#include <type_traits>
#include <unordered_map>
#include <unordered_set>
#include <vector>

class LCFGPlane;
class LCFGBasicBlock;
class LCFGElement {
  friend class LCFG;

public:
  enum Kind {
    Statement,
    Jump,
    JumpTarget,
  };
  //
  LCFGElement(Kind k) : _kind(k) {}
  virtual ~LCFGElement() {}
  virtual void dump(const clang::ASTContext &ctx) = 0;
  Kind kind() { return _kind; }
  auto plane_begin() const { return _planes.begin(); }
  auto plane_end() const { return _planes.end(); }
  const LCFGBasicBlock *basic_block() const { return _basic_block; }

protected:
  void basic_block(LCFGBasicBlock *b) { _basic_block = b; }
  template <typename _It> void add_plane(_It begin, _It end) {
    _planes.insert(_planes.end(), begin, end);
  }
  void id(int i) { _id = i; }

protected:
  Kind _kind;
  int _id;
  LCFGBasicBlock *_basic_block = nullptr;
  std::vector<LCFGPlane *> _planes;
};

class LCFGStatement : public LCFGElement {
public:
  LCFGStatement(clang::Stmt *stmt) : LCFGElement(Statement), _stmt(stmt) {}

  virtual void dump(const clang::ASTContext &ctx) override;
  const clang::Stmt *stmt() const { return _stmt; }

private:
  clang::Stmt *_stmt;
};

struct LCFGJumpTarget : public LCFGElement {
  LCFGJumpTarget() : LCFGElement(JumpTarget) {}
  virtual void dump(const clang::ASTContext &ctx) override {
    llvm::errs() << "JumpTarget\n";
  }
};
struct LCFGJump : public LCFGElement {
  LCFGJump(LCFGJumpTarget *t = nullptr) : LCFGElement(Jump), target(t) {}
  virtual void dump(const clang::ASTContext &ctx) override {
    llvm::errs() << "Jump\n";
  }
  LCFGJumpTarget *target;
};

class LCFGPlane {
  friend class LCFG;

public:
  LCFGPlane(int id) : _id(id) {}
  void dump(const clang::ASTContext &ctx);
  int id() { return _id; }

private:
  int _id;
  void dominator(LCFGElement *e) { _dominator = e; }
  void add_element(LCFGElement *e) { _elements.push_back(e); }

private:
  LCFGElement *_dominator = nullptr;
  std::vector<LCFGElement *> _elements;
};
class LCFGBasicBlock {
  friend class LCFG;

public:
  LCFGBasicBlock(int id) : _id(id) {}
  ~LCFGBasicBlock() {}
  void dump(const clang::ASTContext &ctx);
  int id() const { return _id; }
  auto element_begin() const { return _element_sequence.begin(); }
  auto element_end() const { return _element_sequence.end(); }
  auto succ_begin() const { return _out_blocks.begin(); }
  auto succ_end() const { return _out_blocks.end(); }
  auto pred_begin() const { return _in_blocks.begin(); }
  auto pred_end() const { return _in_blocks.end(); }

private:
  size_t size() { return _element_sequence.size(); }
  void id(int i) { _id = i; }
  void append_element(LCFGElement *e) { _element_sequence.push_back(e); }
  void erase_pred(LCFGBasicBlock *b) { _in_blocks.erase(b); }
  void erase_succ(LCFGBasicBlock *b) { _out_blocks.erase(b); }
  void add_pred(LCFGBasicBlock *b) { _in_blocks.insert(b); }
  void add_succ(LCFGBasicBlock *b) { _out_blocks.insert(b); }

  int _id;
  std::vector<LCFGElement *> _element_sequence;
  std::unordered_set<LCFGBasicBlock *> _in_blocks;
  std::unordered_set<LCFGBasicBlock *> _out_blocks;
};
class LCFG {
private:
  struct BuildContext {
    std::vector<LCFGPlane *> plane_stack;
    std::vector<std::vector<LCFGJump *>> next_jmps_stack;
    LCFGJumpTarget *exit_jmp_target;
  };

public:
  LCFG(clang::FunctionDecl *decl) : _function_decl(decl) { build(); }
  ~LCFG() {
    for (auto i : _basic_blocks)
      delete i;
    for (auto i : _planes)
      delete i;
    for (auto i : _elements)
      delete i;
  }
  void dump() {
    const clang::ASTContext &ctx = _function_decl->getASTContext();
    for (auto i : _basic_blocks) {
      i->dump(ctx);
    }
  }
  auto element_begin() const { return _elements.begin(); }
  auto element_end() const { return _elements.end(); }
  auto basic_block_begin() const { return _basic_blocks.begin(); }
  auto basic_block_end() const { return _basic_blocks.end(); }
  size_t basic_block_size() const { return _basic_blocks.size(); }

private:
  LCFGBasicBlock *new_basic_block() {
    LCFGBasicBlock *blk = new LCFGBasicBlock(_basic_blocks.size());
    _basic_blocks.push_back(blk);
    return blk;
  }
  template <typename _T, typename... _Args>
  typename std::enable_if<std::is_base_of<LCFGElement, _T>::value, _T>::type *
  new_element(_Args &&... args) {
    _T *elm = new _T(std::forward<_Args>(args)...);
    _elements.push_back(elm);
    return elm;
  }
  LCFGPlane *new_plane() {
    LCFGPlane *plane = new LCFGPlane(_planes.size());
    _planes.push_back(plane);
    return plane;
  }
  void bind_element_plane(LCFGElement *elm, std::vector<LCFGPlane *> &planes) {
    elm->add_plane(planes.begin(), planes.end());
    for (auto i : planes) {
      i->add_element(elm);
    }
  }
  void link_basic_block(LCFGBasicBlock *from, LCFGBasicBlock *to) {
    from->add_succ(to);
    to->add_pred(from);
  }
  void unlink_basic_block(LCFGBasicBlock *from, LCFGBasicBlock *to) {
    from->erase_succ(to);
    to->erase_pred(from);
  }
  void add_element_to_basic_block(LCFGElement *e, LCFGBasicBlock *b) {
    b->append_element(e);
    e->basic_block(b);
  }
  void build();
  void dispatch_element_generation(clang::Stmt *stmt, BuildContext &ctx);
  void generate_element(clang::CompoundStmt *stmt, BuildContext &ctx);
  void generate_element(clang::IfStmt *stmt, BuildContext &ctx);
  void generate_element(clang::ForStmt *stmt, BuildContext &ctx);
  void generate_element(clang::WhileStmt *stmt, BuildContext &ctx);
  void generate_element(clang::DoStmt *stmt, BuildContext &ctx);
  void generate_element(clang::ReturnStmt *stmt, BuildContext &ctx);
  void generate_element(clang::ContinueStmt *stmt, BuildContext &ctx);
  void generate_element(clang::BreakStmt *stmt, BuildContext &ctx);
  void generate_element(clang::Stmt *stmt, BuildContext &ctx);
  void remove_useless_element();
  void remove_empty_basic_block();
  void remove_dead_basic_block();
  void remove_illegal_successor();

private:
  clang::FunctionDecl *_function_decl;
  std::vector<LCFGBasicBlock *> _basic_blocks;
  std::vector<LCFGElement *> _elements;
  std::vector<LCFGPlane *> _planes;
  LCFGBasicBlock *_entry_block = nullptr;
  LCFGBasicBlock *_exit_block = nullptr;
  LCFGPlane *_main_plane = nullptr;
};

#endif